from tabulate import tabulate
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from src.norm import TestAlgorithm
from src.const import *
from util import *
pd.set_option('display.max_rows', 100)

def load_data(vid:str, start_time:int=None):
    df = pd.read_csv(RESULTPATH.joinpath(f'{vid}.csv'))
    df['score'] = df['score'].apply(lambda x:int(x[:-1])/100)

    if start_time and 'collection_time' in df.columns:
        df['collection_time'] = df['collection_time'].astype('int')
        df = df[df['collection_time'] >= start_time]  # 时间范围
    
    df1 = df.groupby('indicator').get_group('1 recall')
    df2 = df.groupby('indicator').get_group('1 precision')
    df1, df2 = map(lambda x:x.reset_index(), [df1, df2])
    other_fields = ['support', 'model_name', 'predict_ratio']
    score_df = pd.concat([df1['score'], df2['score'], 
                          df1[other_fields]], axis=1)
    score_df.columns = ['recall', 'precision'] + other_fields
    
    data = score_df['model_name'].str.extract(r'[model_]?(.*?)_([v\.0-9]+)-?.*')
    score_df['model_name'] = data[0]
    score_df['version'] = data[1]  # 分离挂载的字段
    
    return score_df

def vis_scatter(*versions:str, start_time:int=None):
    score_df = pd.DataFrame()
    main_version = versions[0]  # 主版本
    for version in versions:
        df = load_data(version, start_time)
        score_df = pd.concat([score_df, df])
    score_df = score_df[score_df['support'] != '-']
    score_df['support'] = score_df['support'].astype('float')
    for model_name in score_df['model_name'].unique():
        df:pd.DataFrame = score_df[score_df['model_name']==model_name]
        score = TestAlgorithm.weighted_average(scores=df[['recall', 'precision']].to_numpy(), 
                                               weights=df['support'].to_numpy())
        x, y = score.tolist()

        title = model_name
        if 'collection_time' in df.columns:
            ct = df['collection_time'].copy().astype(str)
            title = ct.iloc[0] + '-' + ct.iloc[-1]
        
        fig = px.scatter(df, x='recall', y='precision', color='version', size='support', title=title)
        fig.add_trace(go.Scatter(x=[x], y=[y], showlegend=False, 
                                marker=dict(size=5, color='red', symbol='triangle-up')))
        fig.add_annotation(x=x, y=y, text=f'centre({to_percent(x)}, {to_percent(y)})')
        fig.write_image(str(VISPATH.joinpath(f'{model_name}_scatter_{main_version}.jpg')))

def _my_add_traces(x:list[Any], df:pd.DataFrame, fields:list[str], 
                   fig:go.Figure, name:str, color:Any, 
                   axis:Optional[int]=0, show_text:bool=False):
    kwargs = dict(mode='lines+markers+text', textposition='top center') if show_text else {}
    marker = go.scatter.Marker(color=color)
    for i, field in enumerate(fields):
        if show_text:
            kwargs['text'] = df[field].map(to_percent)
        if i == 0:
            trace = go.Scatter(x=x, y=df[field], name=name, marker=marker, **kwargs)
            fig.add_trace(trace=trace, row=1, col=1)
        else:
            trace = go.Scatter(x=x, y=df[field], showlegend=False, marker=marker, **kwargs)
            if axis == 0 or axis is None:
                fig.add_trace(trace=trace, row=i+1, col=1)
            elif axis == 1:
                fig.add_trace(trace=trace, row=1, col=i+1)

def main_vis(*versions:str, start_time:int=None):
    fields = ['recall', 'precision']
    rows = 2 if len(versions) else 1
    fig = make_subplots(rows=rows, cols=1, shared_xaxes=True, row_titles=fields)
    colors = ['lightgray', 'lightgreen', 'lightblue']
    df = pd.DataFrame()
    main_version = versions[0]  # 主版本
    for version, color in zip(versions, colors):
        score_df = load_data(vid=version, start_time=start_time)
        score_df = score_df[score_df['support'] == '-']
        model_names = score_df['model_name']  # 分组用
        df = pd.concat([df, score_df])
        # score_df['predict_ratio']
        _my_add_traces(x=model_names, df=score_df, fields=fields, 
                       fig=fig, name=version, color=color)
    if len(versions) > 1:  # 多版本
        mean_score_df = df.groupby('model_name')[fields].mean().loc[model_names]
        _my_add_traces(x=model_names, df=mean_score_df, fields=fields, 
                       fig=fig, name='mean', color='red', show_text=True)
    fig.write_image(str(VISPATH.joinpath(f'line_{main_version}.jpg')))
    # fig.show()

def show_table(*versions:str):
    for version in versions:
        print(tabulate(load_data(version), headers='keys', tablefmt='pretty'))

def main():
    # vis_scatter(*['v2', 'v2.1', 'v2.2'])
    # main_vis(*['v2', 'v2.1', 'v2.2'])
    # main_vis('v3')
    show_table('v3-mNl')

if __name__ == '__main__':
    main()
